Skip to content

feat(tilelib): materialize PTODSL in process - #1103

Merged
mouliangyu merged 11 commits into
mainfrom
codex/embedded-ptodsl-context
Aug 7, 2026
Merged

feat(tilelib): materialize PTODSL in process#1103
mouliangyu merged 11 commits into
mainfrom
codex/embedded-ptodsl-context

Conversation

@mouliangyu

@mouliangyu mouliangyu commented Aug 3, 2026

Copy link
Copy Markdown
Collaborator

背景

TileOp 展开此前通过独立 Python daemon 完成 PTODSL 模板选择和 MLIR 生成,需要维护子进程、Unix socket、JSON wire protocol、临时文件以及 MLIR 文本打印/解析。主线移除 legacy TileLang backend 后,这套隔离层只剩 PTODSL 使用,也阻碍了模板实例化直接复用当前编译的 MLIRContext。

主要变更

  • 删除 TileLib daemon manager、Python socket/client/helper/wire 实现、daemon 专用测试和相关 CLI/环境配置。
  • 新增与 Python 无关的 TileLibService C++ 边界:候选查询只传 specialization metadata,模板实例化接收当前 MLIRContext 并返回同 context 下的 source module handle。
  • _core 模块初始化时安装进程级唯一 TileLibRuntime service;runtime 不绑定某个 compilation context,也不在 pass 中保存 Python 对象。
  • 每次 _core.main 创建独立 Python Context owner,native PTOASContext 在该调用的动态范围内借用同一个底层 MLIRContext;连续多次编译共享 runtime,但 context 相互独立。
  • Python callback 通过 MLIR C API capsule 找回当前 context 的 Python wrapper,在显式传入的 context 中 materialize PTODSL module;禁止 compiler path 隐式创建新 context,并校验返回 module 的 context identity。
  • C++ 在 Python module owner 存活期间只借用 source module,按 entry symbol clone 所需函数及 helper 到 caller module;不转移 Python module ownership,也不走 MLIR 文本或 bytecode往返。
  • InsertTemplateAttributes / ExpandTileOp 恢复为 TableGen 标准注册、默认可构造的 pass。pass 从 getContext() 获取当前 context,并从进程级 runtime 取得 service,因此可继续留在连续 pass pipeline,也支持 targeted IR instrumentation。
  • 保留主线的 isTileLibExpandableOp 过滤:没有可展开 TileOp 时 pass 直接返回,不要求 runtime。
  • 将 standalone .mlir_text() 与 compiler materialization API 分开:前者创建 fresh context,后者必须调用 build_module_in_context(context)
  • 将 TileValue 的 valid-shape 派生值改成按需构造的临时 view,避免 surface wrapper 与 tracing/session 之间形成循环引用,无需 weakref 或主动 wrapper invalidation。
  • 更新 TileLib 选择逻辑、CLI/wheel/ST 脚本、设计文档和回归测试。

生命周期与边界

  • Python Context owner 负责底层 MLIRContext 生命周期;PTOASContext 和 pass 只借用当前 compilation context。
  • TileLibRuntime 只管理 process-wide service,module teardown 时通过 capsule 注销,避免跨 DSO/CPython 生命周期留下悬挂全局对象。
  • native compiler 主流程释放 GIL;同步进入 metadata/materialize callback 时 reacquire GIL,不创建额外 worker thread 或 subinterpreter。
  • Python 返回的 module/operation handle 不跨 context 使用;C++ clone 完成后即可释放 Python source module wrapper。
  • 本 PR 回归覆盖同一 Python 进程内的顺序重复 _core.main。并发 driver invocation 仍受 LLVM 全局 command-line state 限制,不在本 PR 承诺范围内。

验证

  • cmake --build .work/tileop-embedded-python-study/build --target PTOASPythonPackage -j4
  • ctest --test-dir .work/tileop-embedded-python-study/build --output-on-failure -L PTODSL -j4:25/25 通过
  • llvm-lit -sv:以下 4 个 TileOp/VPTO 用例全部通过
    • expand_tile_op_ptodsl_tadd.pto
    • expand_tile_op_ptodsl_tsub.pto
    • expand_tile_op_ptodsl_view_stride_cache.pto
    • vpto_pipeline_vmi_after_tileop_expand.pto
  • 新增 test_ptoas_runtime.py,验证同一进程连续两次 _core.main 使用不同 compilation context 且共享 process runtime,并均完成 TileOp 展开。
  • python3 -m py_compilegit diff --check 通过。

构建和测试使用 LLVM/MLIR 21.1.8。

@mouliangyu
mouliangyu force-pushed the codex/embedded-ptodsl-context branch from 279b47f to 7b080c2 Compare August 4, 2026 05:43
@mouliangyu
mouliangyu marked this pull request as ready for review August 7, 2026 09:15
@mouliangyu
mouliangyu merged commit 988d50e into main Aug 7, 2026
16 of 19 checks passed
@reedhecre

Copy link
Copy Markdown

A5 板测失败

  • 触发方式:merged
  • 源码提交:988d50e24521
  • 结果汇总:OK 20 / FAIL 1 / SKIP 0
  • 日志:/root/ptoas-board-monitor-a5/logs/20260807_174706_merged_pr1103.log
  • 失败阶段:board-validation-qwen / exit=1

失败用例

  • rope_kv_cache (run, exit=2)

@reedhecre

Copy link
Copy Markdown

A3 板测完成(有跳过)

  • 触发方式:merged
  • 源码提交:988d50e24521
  • 结果汇总:OK 220 / FAIL 0 / SKIP 3
  • 日志:/home/zhongxuan/ptoas-board-monitor/runtime/logs/20260807_025809_merged_pr1103.log
  • 结果 TSV:/home/zhongxuan/ptoas-board-monitor/runtime/logs/20260807_025809_merged_pr1103.tsv

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants